package handlers

import (
  "context"
  "crypto/tls"
  "github.com/bytedance/sonic"
  "github.com/cloudwego/hertz/pkg/app"
  "github.com/cloudwego/hertz/pkg/app/client"
  "github.com/cloudwego/hertz/pkg/common/hlog"
  "github.com/cloudwego/hertz/pkg/protocol"
  "github.com/hertz-contrib/sse"
  "go-llm-proxy/hutils"
  "strings"
)

var OpenAiApiPrefix = "https://api.openai.com/v1"
var OpenAiChatCompletionUrl = OpenAiApiPrefix + "/chat/completions"
var OpenAiModelUrl = OpenAiApiPrefix + "/models"

func OpenAiModels(ctx context.Context, reqCtx *app.RequestContext) {
  uri := reqCtx.URI()
  hlog.Info("uri:", uri.String())
  headers := make(map[string]string)
  reqCtx.Request.Header.VisitAll(func(key, value []byte) {
    headers[string(key)] = string(value)
  })
  //headers.put("host", "api.openai.com");
  headers["Host"] = "api.openai.com"

  // 配置 TLS
  tlsConfig := &tls.Config{
    InsecureSkipVerify: false, // 确保验证服务器证书
  }

  // 创建支持 TLS 的 HTTP 客户端
  httpClient, err := client.NewClient(client.WithTLSConfig(tlsConfig))
  if err != nil {
    hlog.Error("Failed to create HTTP client:", err)
    hutils.Fail500(reqCtx, err)
    return
  }

  request := &protocol.Request{}
  response := &protocol.Response{}
  request.SetRequestURI(OpenAiModelUrl)
  request.SetMethod("GET")
  request.SetHeaders(headers)

  // 执行请求
  err = httpClient.Do(context.Background(), request, response)
  defer response.CloseBodyStream()
  if err != nil {
    hlog.Error("url:", OpenAiApiPrefix, " error:", err.Error())
    hutils.Fail500(reqCtx, err)
    return
  }
  // 设置响应头和状态码
  response.Header.VisitAll(func(key, value []byte) {
    reqCtx.Response.Header.Set(string(key), string(value))
  })
  // 设置响应体流
  reqCtx.Response.SetBody(response.Body())
}
func OpenAiV1ChatCompletions(ctx context.Context, reqCtx *app.RequestContext) {
  // Read request body and headers
  body, _ := reqCtx.Body()
  // Decode the body into a map
  var requestMap map[string]interface{}
  sonic.ConfigDefault.Unmarshal(body, &requestMap)

  headers := make(map[string]string)
  reqCtx.Request.Header.VisitAll(func(key, value []byte) {
    headers[string(key)] = string(value)
  })
  //headers.put("host", "api.openai.com");
  headers["Host"] = "api.openai.com"
  hlog.Info("body:", string(body))
  hlog.Info("headers:", headers)
  if requestMap["stream"] == true {
    // Setup client to connect to the remote server
    client := sse.NewClient(OpenAiChatCompletionUrl)
    client.SetMethod("POST")
    client.SetHeaders(headers)
    client.SetBody(body)

    // Setup the stream for the original client
    var sEvent *sse.Stream = sse.NewStream(reqCtx)
    errChan := make(chan error)
    var completeContent strings.Builder
    go func() {
      err := client.Subscribe(func(msg *sse.Event) {
        if msg.Data != nil {
          // Forwarding the received event back to the original client
          event := &sse.Event{
            Data: msg.Data,
          }
          err := sEvent.Publish(event)
          if err != nil {
            hlog.CtxErrorf(ctx, "failed to send event to client: %sEvent", err)
            return
          }
          printResponseContent(msg, &completeContent)
        }
      })
      errChan <- err
    }()

    select {
    case err := <-errChan:
      if err != nil {
        hlog.CtxErrorf(ctx, "error from remote server: %sEvent", err)
        hutils.Fail500(reqCtx, err)
        return
      }
    }
  } else {
    httpClient, _ := client.NewClient(client.WithTLSConfig(&tls.Config{
      InsecureSkipVerify: true, //InsecureSkipVerify: true 会跳过对服务器证书的验证，这在生产环境中是不推荐的
    }))
    request := &protocol.Request{}
    response := &protocol.Response{}
    request.SetRequestURI(OpenAiChatCompletionUrl)
    request.SetMethod("POST")
    request.SetHeaders(headers)
    request.SetBody(body)

    // 执行请求
    var err = httpClient.Do(context.Background(), request, response)
    defer response.CloseBodyStream()
    if err != nil {
      hlog.Error("error:", err)
      hutils.Fail500(reqCtx, err)
      return
    }
    // 设置响应头和状态码
    response.Header.VisitAll(func(key, value []byte) {
      reqCtx.Response.Header.Set(string(key), string(value))
    })
    // 设置响应体流
    reqCtx.Response.SetBody(response.Body())
  }
}
